﻿using System;
using System.Collections.Generic;
using System.Text;
using System.Net.Sockets;
using System.Net;
using System.IO;
using ReliableSession;
using ReliableSession.src;
using log4net;
using System.Diagnostics;
using System.Threading;
using seqType = System.UInt64;  /// Type for sequence numbers

namespace ReliableSession
{
    /// <summary>
    /// A chunk 
    /// </summary>
    public class Chunk
    {
        /// <summary>
        /// Start position of this chunk.  Reduces the number of additions when comparing chunks.
        /// </summary>
        public seqType m_startPos;
        public seqType m_endPos;
        public ushort m_length;

        public Chunk(seqType endPos, ushort length)
        {
            m_startPos = endPos - length;
            m_endPos = endPos;
            m_length = length;
        }
    }

    /// <summary>
    /// Creates a reliable UDP transport layer
    /// </summary>
    public class Transport
    {
        private static readonly ILog log = LogManager.GetLogger(typeof(Transport));
        /// <summary>
        /// Receive data buffer
        /// </summary>
        /// 
        private FifoBuffer m_rxFifo = new FifoBuffer(64 * 1024);
        private byte[] m_rxBuffer = new byte[64 * 1024];

        /// <summary>
        /// Send data buffer.
        /// </summary>
        private byte[] m_txBuffer = new byte[64 * 1024];

        /// <summary>
        /// A copy of the data being sent for potential re-send.
        /// </summary>
        private FifoBuffer m_txFifo = new FifoBuffer(64 * 1024);

        private EndPoint m_remoteAddress;
        public EndPoint m_RemoteAddress
        {
            get { return m_remoteAddress; }
            set { m_remoteAddress = value; }
        }

        /// <summary>
        /// true if this is opened as a server socket.
        /// </summary>
        private bool m_isServerSocket = false;

        /// <summary>
        /// The actual underlying socket
        /// </summary>
        private Socket m_socket;
        private BinaryReader m_rxReader;
        //private BinaryReader m_rxReaderCopy;

        private BinaryWriter m_txWriter;

        private TransportHeader m_header;
        /// <summary>
        /// header to send the ack back to the peer.  
        /// TODO: Remove this
        /// </summary>
        private TransportHeader m_ackHeader;

        private seqType m_remoteSeqNumber;
        private seqType m_mySeqNumber;
        /// <summary>
        /// The sequence number last acknowledged by the peer
        /// </summary>
        private seqType m_lastAckedSeqNumber;
        private bool m_handshakeDone = false;

        /// <summary>
        /// The maximum length of a single packet which is negotiated at startup
        /// and may be renegotiated.
        /// </summary>
        private ushort m_maxPacketSize = 512;

        /// <summary>
        /// The maximum number of packets that the sender can send before it waits for an ack from the peer.
        /// </summary>
        private int m_maxSendPacketsBeforeAck = 10;

        /// <summary>
        /// Number of packets sent after the last time the peer acknowledged the receipt
        /// </summary>
        private int m_packetsSentAfterLastAckFromPeer = 0;

        /// <summary>
        /// An ACK is pending to be sent.  Send it ASAP.
        /// </summary>
        private bool m_sendAck = false;
        private bool m_readBufferFull;
        private int m_sleepTimeBetweenAckReads = 5;
        public TransportStatistics m_statistics { get; set; }

        public ushort m_headerLength = 0;
        private string m_debugStr = "";
        private bool m_trace = false;
        public bool m_debug = false;
        private bool m_forceAck;

        private List<Chunk> m_receivedChunks = new List<Chunk>();
        public Transport()
        {
            m_remoteAddress = new IPEndPoint(IPAddress.Any, 0);
            m_rxReader = new BinaryReader(new MemoryStream(m_rxBuffer));
            m_txWriter = new BinaryWriter(new MemoryStream(m_txBuffer));
            m_lastAckedSeqNumber = m_mySeqNumber = 0;
            m_remoteSeqNumber = 0;
            m_header = new TransportHeader();
            m_ackHeader = new TransportHeader();
            m_statistics = new TransportStatistics();

            // calculate the header length
            m_header.serialize(m_txWriter);
            m_headerLength = (ushort)m_txWriter.BaseStream.Position;
            m_txWriter.BaseStream.Position = 0;
            m_maxPacketSize -= m_headerLength;
        }

        private void resetOutputBuffer()
        {
            m_txWriter.BaseStream.Position = 0;   // reset the write position
        }

        /// <summary>
        /// Open a server socket at the specified port.
        /// Throws the underlying socket exception if there is an error.
        /// </summary>
        /// <param name="port"></param>
        public void openServerSocket(int port)
        {
            lock (this)
            {
                if (m_socket != null && (m_socket.Connected == false || m_socket.IsBound == false))
                {
                    throw new SocketException((int)SocketError.IsConnected);
                }
                EndPoint endPoint = new IPEndPoint(IPAddress.Any, port);
                m_socket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
                m_socket.Bind(endPoint);
                m_isServerSocket = true;
            }
        }

        /// <summary>
        /// Accept an incoming socket request and return a newly created client socket
        /// </summary>
        /// <returns></returns>
        public Transport accept()
        {
            lock (this)
            {
                if (m_socket == null || m_socket.IsBound == false)
                {
                    // the socket must be open.
                    throw new SocketException((int)SocketError.NotConnected);
                }
                if (!m_isServerSocket)
                {
                    // the socket must be opened as a server socket.
                    throw new SocketException((int)SocketError.OperationNotSupported);
                }
                // wait for data on the socket.
                Transport transport = new Transport();
                transport.serverSideInitialHandshake(m_socket);
                return transport;
            }
        }

        /// <summary>
        /// Perform the server side initial handshake.
        /// 1. Receive SYN from client
        /// 2. Send SYN/ACK to client
        /// 3. Read SYN/ACK from client
        /// </summary>
        private void serverSideInitialHandshake(Socket socket)
        {
            TransportHeader header = m_header;
            {
                // STEP 2
                // Read the first packet from the client to start establishing the 3 way handshake
                deserializePacket(socket,false);
                /// Read the first SYN from the client and send back an ACK

                if (header.isBitSet(TransportHeader.BitFlags.SYN))
                {
                    //log.Debug("Server got SYN from client: " + info());
                    // if the synchronization bit is set then open a local socket and send the ack
                    openClientSocketAndBind();     // open a local port.
                    //m_socket.ReceiveTimeout = 5 * 1000; // default timeout during handshake.
                    m_remoteSeqNumber = header.m_seqNumber;
                    // send an ACK back to the client.
                    serialize("S2C SYN/ACK",m_header, true, true, false,null, 0, 0);
                }
                else
                {
                    resetConnection("Server socket did not get SYN from client");
                }
            }

            {
                // STEP 3
                // read the syn,ack from the sender
                header.clearBitFlag();
                deserializePacket(m_socket,false);
                bool synAndAck = header.isBitSet(TransportHeader.BitFlags.SYN) && header.isBitSet(TransportHeader.BitFlags.ACK);
                if (!synAndAck)
                {
                    // if the other side did not send a syn/ack then close the connection and throw an exception
                    resetConnection("Server socket did not get SYN and/or ACK from client");
                }
                else if (header.m_ackNumber != m_mySeqNumber)
                {
                    resetConnection("Incorrect seqn number from client, expecting <" + m_mySeqNumber + "> got <" + header.m_ackNumber + ">");
                }
                else
                {
                    m_handshakeDone = true;
                    m_statistics.m_openedTime = DateTime.Now;
                    //log.Debug("Server got SYN and ACK: " + info());
                }
            }
        }

        #region Client side 3 way handshake
        /// <summary>
        /// Open a Reliable Datagram client socket to the specified server@port
        /// </summary>
        /// <param name="ipaddress"></param>
        /// <param name="port"></param>
        public void openClientSocket(string ipaddress, int port)
        {
            openClientSocketAndBind();
            {
                // Send SYN to the server                
                m_lastAckedSeqNumber = m_mySeqNumber = 0;    // set the initial sequence number
                m_RemoteAddress = new IPEndPoint(IPAddress.Parse(ipaddress), port);
                serialize("C2S SYN",m_header, false, true, false,null, 0, 0,false,true);
            }
         
            {
                // send SYN/ACK to server
                this.deserializePacket(m_socket,false);
                if (!m_header.isBitSet(TransportHeader.BitFlags.SYN))
                {
                    resetConnection("Client socket did not get SYN from server");
                }
                m_remoteSeqNumber = m_header.m_seqNumber;   // remember the peer's sequence number
                serialize("C2S SYN/ACK", m_header,true, true, false,null, 0, 0,false,true);
                m_handshakeDone = true;
                m_statistics.m_openedTime = DateTime.Now;
            }
        }
        #endregion


        /// <summary>
        /// Deserialize the data that is available in the m_rxBuffer
        /// <returns>The number of bytes read from the body of the packet (does not include the header)</returns>
        /// </summary>
        private int deserializePacket(Socket socket,bool peekMode)
        {
            TransportHeader header = m_header;
            int received = -1;
            do
            {
                // reset the stream pointer
                m_rxReader.BaseStream.Position = 0;
                if (peekMode)
                {
                    if ((received = socket.Available) <= 0)
                    {
                        // nothing to do
                        break;
                    }
                }
                received = socket.ReceiveFrom(m_rxBuffer, m_rxBuffer.Length, SocketFlags.None, ref m_remoteAddress);

                int receivedLen = received;
                // read the header part
                header.deserialize(m_rxReader);
                int len = header.m_dataLength;
                if (header.m_dataLength <= m_maxPacketSize)
                {
                    ushort index = 0;
                    // ignore sequence check if handshake is not complete yet.
                    received = m_rxReader.Read(m_rxBuffer, 0, len);  // read the body
                    receivedLen = received;
                    Debug.Assert(received == len);
                    if (header.isBitSet(TransportHeader.BitFlags.ACK))
                    {
                        // got an acknowledgement from the peer
                        bool ackInRange =
                            TransportMath.gte(m_lastAckedSeqNumber, header.m_ackNumber) &&
                            TransportMath.lte(m_mySeqNumber, header.m_ackNumber);
                        if (ackInRange)
                        {
                            if (TransportMath.gt(m_lastAckedSeqNumber, header.m_ackNumber))
                            {
                                // compact the send buffer to the length of the last acked sequence number.
                                ushort compactLen = (ushort)(header.m_ackNumber - m_lastAckedSeqNumber);
                                m_txFifo.compact(compactLen);
                                m_statistics.incrementAcks();
                                m_lastAckedSeqNumber = header.m_ackNumber;
                                if (m_debug)
                                {
                                    log.Debug("Got ack=" + m_lastAckedSeqNumber + " MySeq=" + m_mySeqNumber + " Compact=" + compactLen + " Available=" + m_txFifo.dataAvailable() + "ReadPos=" + m_txFifo.readPosition() + "Writepos=" + m_txFifo.writePosition());
                                }
                            }
                        }
                        else
                        {
                            // drop this packet
                            log.Debug("Peer acked a sequence number that is not in sequence. acked=" + header.m_ackNumber + " Our last ack " + m_lastAckedSeqNumber);
                            break;
                        }
                        m_packetsSentAfterLastAckFromPeer = 0;
                        if (len != 0)
                        {
                            m_sendAck = true;
                        }
                    }

                    if (header.isBitSet(TransportHeader.BitFlags.RST))
                    {
                        m_socket.Close();
                    }

                    bool greater = TransportMath.gt(m_remoteSeqNumber, header.m_seqNumber);
                    if (greater)
                    {
                        // find the amount of space left in the read copy buffer
                        long remaining = m_rxFifo.spaceAvailable();
                        index = 0;
                        ushort diff = (ushort)(header.m_seqNumber - m_remoteSeqNumber);
                        if (len > remaining)
                        {
                            // if the incoming data cannot be accomodated in the copy buffer then we have to tell the
                            // other side to wait, or atleast not send an ack that we read the current data.
                            m_readBufferFull = true;
                            break;
                        }

                        if (diff < len)
                        {
                            // there is some data that we have already seen, discard it.
                            index = (ushort)(len - diff);
                            len = diff;
                            received = diff;
                        }
                        else if (diff > len)
                        {
                            // some packet got lost
                            m_statistics.incrementLost(diff - len);
                            m_sendAck = true;
                            // write the data at an offset
                            int startPos = diff;
                            m_rxFifo.write(m_rxBuffer, startPos, index, len);
                            addIndexToList(header.m_seqNumber, len);
                        } 
                        else 
                        {
                            // copy the data received into the backup buffer.
                            m_rxFifo.write(m_rxBuffer, index, len);
                            m_readBufferFull = false;
                            m_remoteSeqNumber = m_header.m_seqNumber;
                            m_statistics.incrementReceived(len);
                            return received;
                        }
                    }
                }
                else
                {
                    // we have to drop this packet
                    // TODO: Send back an error saying the packet is too big.
                    log.Error("Packet is bigger than the max size");
                }
            } while (socket.Available > 0);
            return -1;
        }

        private void addIndexToList(seqType inSeq, int length)
        {
            // adding to index does not work with datatype other than 64 bit because rollover is not considered.
            // Even in the case of 64 bit it is assumed that the sequence number started at a low value.
            Debug.Assert(typeof(seqType) == typeof(System.UInt64));

            ushort len = (ushort)length;
            if (m_receivedChunks.Count == 0)
            {
                Chunk chunk = new Chunk(inSeq, len);
                m_receivedChunks.Add(chunk);
            }
            else
            {
                bool inserted = false;
                for(int i=0; i <m_receivedChunks.Count; i++) 
                {
                    Chunk chunk = m_receivedChunks[i];
                    // the incoming chunk must be placed before the current chunk or coalesced with it.
                    if (inSeq == chunk.m_startPos)
                    {
                        chunk.m_startPos = inSeq - len;
                        chunk.m_length += len;
                        inserted = true;
                        if (i -1 >= 0)
                        {
                            Chunk prevChunk = m_receivedChunks[i-1];
                            if (prevChunk.m_endPos == chunk.m_startPos)
                            {
                                prevChunk.m_endPos = chunk.m_endPos;
                                prevChunk.m_length += chunk.m_length;
                                m_receivedChunks.RemoveAt(i);
                            }
                        }
                        break;
                    }
                    else if(chunk.m_endPos > inSeq)
                    {
                        // the new chunk must be inserted before the current one.
                        Chunk newChunk = new Chunk(inSeq, len);
                        m_receivedChunks.Insert(i, newChunk);
                        inserted = true;
                        break;
                    }
                }
                if (!inserted)
                {
                    Chunk chunk = new Chunk(inSeq, len);
                    m_receivedChunks.Add(chunk);
                }
            }
        }

        private void resetConnection(string msg)
        {
            closeConnection();
            log.Debug(msg);
            throw new SocketException((int)SocketError.ConnectionReset);
        }

        /// <summary>
        /// Send the reset packet and close the socket connection.
        /// </summary>
        private void closeConnection()
        {
            if (m_socket != null)
            {
                m_socket.Close();
            }
        }

        /// <summary>
        /// Open a new client socket and bind it to a local random port
        /// </summary>
        private void openClientSocketAndBind()
        {
            if (m_socket != null && (m_socket.Connected == false || m_socket.IsBound == false))
            {
                throw new SocketException((int)SocketError.IsConnected);
            }

            if (m_isServerSocket)
            {
                // cannot convert a server socket into a client socket.
                throw new SocketException((int)SocketError.InvalidArgument);
            }
            // bind to a local port
            m_socket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
            EndPoint ep = new IPEndPoint(IPAddress.Any,0);
            m_socket.Bind(ep);
        }

        private void serialize(String logMsg, TransportHeader header, bool ack, bool syn, bool reset, byte[] data, int pos, int len)
        {
            serialize(logMsg, header, ack, syn,false, data, pos, len, true,true);
        }

        private void serialize(String logMsg,TransportHeader header,bool ack,bool syn,bool reset,byte[] data,int pos,int len,bool makeCopy,bool readAck)
        {
            resetOutputBuffer();
            header.clearBitFlag();

            synchronize(readAck);
            string extra = "";

            if(ack)
            {
                header.setBit(TransportHeader.BitFlags.ACK);
                header.m_ackNumber = m_remoteSeqNumber;
            }
            if (reset)
            {
                header.setBit(TransportHeader.BitFlags.RST);
                header.m_ackNumber = m_remoteSeqNumber;
            }

            if (m_handshakeDone==false && syn)
            {
                // allow syn bit to be sent only if handshake is not complete
                header.setBit(TransportHeader.BitFlags.SYN);
            }

            if (len > m_maxPacketSize)
            {
                len = m_maxPacketSize;
            }
            // prepare the header to be written to the wire
            ushort increment = (ushort)((data == null || len == 0) ? 0 : len);
            if (makeCopy)
            {
                m_mySeqNumber += increment; // increment the sequence number
                header.m_seqNumber = m_mySeqNumber;
            }
            else
            {
                header.m_seqNumber = (ushort)(m_lastAckedSeqNumber + increment);
            }
            if (m_trace)
            {
                extra += "MySeq=" + m_mySeqNumber + " ";
                extra += "  PeerAck=" + header.m_ackNumber + " ";
            }
            header.m_dataLength = increment;
            header.serialize(m_txWriter);
            // prepare for data to be on the wire.
            if (increment > 0)
            {
                m_txWriter.BaseStream.Write(data, pos, increment);
                bool resend = false;
                if (makeCopy)
                {
                    int l = 0;
                    while((l = m_txFifo.write(data, pos, increment)) <= 0 ) {
                        synchronize(readAck);
                    }
                    if (l < increment)
                    {
                        increment = (ushort)l;
                    }
                    m_statistics.incrementSent(increment);
                }
                else
                {
                    resend = true;
                    m_statistics.incrementResent(increment);
                }

                //String s = (resend ? "Resend: " : "Send: ")  + Encoding.ASCII.GetString(m_txBuffer, (int)(m_headerLength), (int)(m_txWriter.BaseStream.Position - m_headerLength));
                //System.Console.WriteLine(s);

                if (m_trace)
                {
                    m_debugStr += resend ? "\nResend: " : "\nSend: ";
                    m_debugStr += " MySeq=" + m_mySeqNumber; // +" LastAck=" + m_lastAckedSeqNumber;
                    m_debugStr += " " + Encoding.ASCII.GetString(m_txBuffer, (int)(m_headerLength), (int)(m_txWriter.BaseStream.Position - m_headerLength)) + "\n";
                }
            }

            m_socket.SendTo(m_txBuffer, (int)(0), (int)m_txWriter.BaseStream.Position, SocketFlags.None, m_remoteAddress);


            if (increment > 0)
            {
                ++m_packetsSentAfterLastAckFromPeer;
            }
            if (logMsg != null && logMsg.Length > 0)
            {
                if (m_debug)
                {
                    log.Debug(logMsg + ": " + extra + " " + header.ToString() + ": Local=" + m_socket.LocalEndPoint.ToString() + " Remote=" + m_remoteAddress.ToString());
                }
            }

            synchronize(readAck);
        }

        private void synchronize(bool readAck)
        {
            if (readAck && timeForAck())
            {
                int numEmptyReads = 0;
                int timeSpentInEmptyReads = 0;
                int numResends = 0;
                bool peekMode = true;
                seqType lastAckedSeqNumber = m_lastAckedSeqNumber;
                int maxEmptyReads = 10;
                int sleepTime = m_sleepTimeBetweenAckReads;
                if (TransportMath.lt(m_lastAckedSeqNumber, m_mySeqNumber))
                {
                    while (TransportMath.lt(m_lastAckedSeqNumber, m_mySeqNumber))
                    {
                        int received = readDataOrAck(null, 0, 0, peekMode);
                        if (lastAckedSeqNumber == m_lastAckedSeqNumber)
                        {
                            // the sequence number did not move
                            ++numEmptyReads;
                            Thread.Sleep(5);

                            if (numEmptyReads >= maxEmptyReads)
                            {
                                // make 5 attempts before waiting for a short period and turning back socket read wait on
                                ++numResends;
                                m_sendAck = true;
                                resendDataOrAck(false);
                                numEmptyReads = 0;
                                Thread.Sleep(sleepTime);
                                sleepTime *= 2;
                                if (sleepTime > 5 * 1000)
                                {
                                    sleepTime = 5 * 1000;
                                }
                                maxEmptyReads *= 2;
                                if (maxEmptyReads >= 10)
                                {
                                    maxEmptyReads = 10;
                                }
                                timeSpentInEmptyReads += m_sleepTimeBetweenAckReads;
                            }
                        }
                        else
                        {
                            lastAckedSeqNumber = m_lastAckedSeqNumber;
                            // we got some data and we are expecting more to be there.
                            peekMode = true;
                            numEmptyReads = 0;
                            timeSpentInEmptyReads = 0;
                            sleepTime = m_sleepTimeBetweenAckReads;
                            maxEmptyReads = 1;
                        }
                    }
                    //Thread.Sleep(1000);
                }
            }
        }

        private bool timeForAck()
        {
            bool value = m_handshakeDone && (m_packetsSentAfterLastAckFromPeer > m_maxSendPacketsBeforeAck) ||
                (DateTime.Now.Ticks - m_statistics.m_lastRecvTime.Ticks < 3 * TimeSpan.TicksPerSecond) ||
                m_forceAck;
            return value;
        }

        public int send(String str)
        {
            byte[] data = Encoding.ASCII.GetBytes(str);
            return this.send(data, 0, data.Length);
        }

        public int send(byte[] data, int pos, int length)
        {
            int sent = 0;
            int remaining = length;
            ushort max = m_maxPacketSize;
            ushort chunk = 0;
            do
            {
                chunk = length > max ? max : (ushort)length;
                serialize("Send",m_header, true, false, false,data, pos + sent, chunk);
                remaining -= chunk;
                sent += chunk;
            } while (remaining > 0);

            return sent;
        }

        public int read(byte[] data, int pos, int length)
        {
            int len = 0;
            while (len == 0)
            {
                len = readDataOrAck(data, pos, length, false);
                resendDataOrAck(true);
            }
            return len;
        }

        private int readDataOrAck(byte[] data, int pos, int length,bool peekMode)
        {
            int len = 0;
            int available = m_rxFifo.dataAvailable();
            if (data != null && available > 0)
            {
                // copy any remaining data from the copy buffer
                len = available > length ? length : available;
                m_rxFifo.read(data, pos, len);
                m_rxFifo.compact();
                return len;
            }
            len = deserializePacket(m_socket,peekMode);
            if (data != null && len >0)
            {
                if (len > length)
                {
                    len = length;
                }
                m_rxFifo.read(data, pos, len);
                m_rxFifo.compact();
            }
            if (!(peekMode == false && len <= 0))
            {
                if (m_debug)
                {
                    log.Debug("Read:" + " PeerSeq=" + m_header.m_seqNumber + " PeerAck=" + m_header.m_ackNumber + " Len=" + len + " Header=" + System.Convert.ToString(m_header.m_bitFlags) + ": Local=" + m_socket.LocalEndPoint.ToString() + " Remote=" + m_remoteAddress.ToString());
                }
            }
            return len;
        }

        /// <summary>
        /// Resend data+ack or just send the ack.
        /// </summary>
        private void resendDataOrAck(bool readAck)
        {
            if (m_sendAck)
            {
                m_ackHeader.clearBitFlag();
                m_ackHeader.m_ackNumber = m_header.m_seqNumber;
                m_sendAck = false;
                int len = (ushort)(m_mySeqNumber - m_lastAckedSeqNumber);
                Debug.Assert(len>=0);
                if (len > 0)
                {
                    int dataAvailable = m_txFifo.dataAvailable();
                    Debug.Assert(len == dataAvailable);
                    if(len > m_maxPacketSize) 
                    {
                        len = m_maxPacketSize;
                    }
                    int index = m_txFifo.readPosition();
                    serialize("Resending data", m_ackHeader, true, false,false, m_txFifo.m_data, index, len, false, readAck);
                }
                else
                {
                    //serialize("Sending back an ACK",m_header, true, false, null, 0, 0);
                    serialize("Sending back ACK", m_header, true, false, false,null, 0, 0);
                }
            }
        }

        public void close()
        {
            m_forceAck = true;
            synchronize(true);
            serialize("Closing", m_header, false, false,true, null, 0, 0);
            m_socket.Close();
        }

        
        public bool testAddToIndex()
        {
            seqType pos = 1;
            ushort length = 0;

            length = 10;
            pos += length;
            this.addIndexToList(pos,length);
            ushort length2 = 5;
            pos += length2;

            seqType pos2 = pos;

            length = 7;
            pos += length;
            this.addIndexToList(pos, length);
            this.addIndexToList(pos2, length2);
            return true;
        }
    }
}
